package org.finesys.chat.core.base.document.splitter;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

import org.finesys.chat.core.base.document.Document;
import org.finesys.chat.core.base.document.Metadata;
import org.finesys.chat.core.base.document.tokenizer.Tokenizer;
import org.finesys.chat.core.base.segment.TextSegment;
import org.finesys.common.core.util.Utils;
import org.finesys.common.core.util.ValidationUtil;

public abstract class HierarchicalDocumentSplitter implements DocumentSplitter {

    private HierarchicalDocumentSplitter overlapSentenceSplitter;
    protected final int maxSegmentSize;
    protected final int maxOverlapSize;
    protected final Tokenizer tokenizer;
    protected final DocumentSplitter subSplitter;

    private HierarchicalDocumentSplitter getOverlapSentenceSplitter() {
        if (this.overlapSentenceSplitter == null) {
            this.overlapSentenceSplitter = new DocumentBySentenceSplitter(1, 0, (Tokenizer) null, (DocumentSplitter) null);
        }

        return this.overlapSentenceSplitter;
    }

    protected HierarchicalDocumentSplitter(int maxSegmentSizeInChars, int maxOverlapSizeInChars) {
        this(maxSegmentSizeInChars, maxOverlapSizeInChars, (Tokenizer) null, (DocumentSplitter) null);
    }

    protected HierarchicalDocumentSplitter(int maxSegmentSizeInChars, int maxOverlapSizeInChars, HierarchicalDocumentSplitter subSplitter) {
        this(maxSegmentSizeInChars, maxOverlapSizeInChars, (Tokenizer) null, subSplitter);
    }

    protected HierarchicalDocumentSplitter(int maxSegmentSizeInTokens, int maxOverlapSizeInTokens, Tokenizer tokenizer) {
        this(maxSegmentSizeInTokens, maxOverlapSizeInTokens, tokenizer, (DocumentSplitter) null);
    }

    protected HierarchicalDocumentSplitter(int maxSegmentSizeInTokens, int maxOverlapSizeInTokens, Tokenizer tokenizer, DocumentSplitter subSplitter) {
        this.maxSegmentSize = ValidationUtil.ensureGreaterThanZero(maxSegmentSizeInTokens, "maxSegmentSize");
        this.maxOverlapSize = ValidationUtil.ensureBetween(maxOverlapSizeInTokens, 0, this.maxSegmentSize, "maxOverlapSize");
        this.tokenizer = tokenizer;
        this.subSplitter = subSplitter == null ? this.defaultSubSplitter() : subSplitter;
    }

    protected abstract String[] split(String var1);

    protected abstract String joinDelimiter();

    protected abstract DocumentSplitter defaultSubSplitter();

    public List<TextSegment> split(Document document) {
        ValidationUtil.ensureNotNull(document, "document");
        List<TextSegment> segments = new ArrayList<>();
        SegmentBuilder segmentBuilder = new SegmentBuilder(this.maxSegmentSize, this::estimateSize, this.joinDelimiter());
        AtomicInteger index = new AtomicInteger(0);
        String[] parts = this.split(document.getText());
        String overlap = null;

        for (String part : parts) {
            int partSize = segmentBuilder.sizeOf(part);
            if (segmentBuilder.hasSpaceFor(partSize)) {
                segmentBuilder.append(part);
            } else {
                if (segmentBuilder.isNotEmpty()) {
                    String segmentText = segmentBuilder.toString();
                    if (!segmentText.equals(overlap)) {
                        segments.add(createSegment(segmentText, document, index.getAndIncrement()));
                        overlap = this.overlapFrom(segmentText);
                        segmentBuilder.reset();
                        segmentBuilder.append(overlap);
                        if (segmentBuilder.hasSpaceFor(partSize)) {
                            segmentBuilder.append(part);
                            continue;
                        }
                    }
                }

                if (this.subSplitter == null) {
                    throw new IllegalArgumentException(String.format(
                            "The text \"%s...\" (%s %s) doesn't fit into maximum segment size (%s %s)",
                            Utils.firstChars(part, 30),
                            this.estimateSize(part),
                            this.tokenizer == null ? "characters" : "tokens",
                            this.maxSegmentSize,
                            this.tokenizer == null ? "characters" : "tokens"));
                }

                segmentBuilder.append(part);
                Iterator<TextSegment> subParts = this.subSplitter.split(Document.from(segmentBuilder.toString())).iterator();

                while (subParts.hasNext()) {
                    TextSegment segment = subParts.next();
                    segments.add(createSegment(segment.getText(), document, index.getAndIncrement()));
                }

                if (!segments.isEmpty()) {
                    TextSegment lastSegment = segments.get(segments.size() - 1);
                    overlap = this.overlapFrom(lastSegment.getText());
                    segmentBuilder.reset();
                    segmentBuilder.append(overlap);
                }
            }
        }

        if (segmentBuilder.isNotEmpty() && !segmentBuilder.toString().equals(overlap)) {
            segments.add(createSegment(segmentBuilder.toString(), document, index.getAndIncrement()));
        }

        return segments;
    }

    String overlapFrom(String segmentText) {
        if (this.maxOverlapSize == 0) {
            return "";
        } else {
            List<String> sentences = Arrays.asList(this.getOverlapSentenceSplitter().split(segmentText));
            Collections.reverse(sentences);
            SegmentBuilder overlapBuilder = new SegmentBuilder(this.maxOverlapSize, this::estimateSize, this.joinDelimiter());
            Iterator<String> sentencesIterator = sentences.iterator();

            while (sentencesIterator.hasNext()) {
                String sentence = (String) sentencesIterator.next();
                if (!overlapBuilder.hasSpaceFor(sentence)) {
                    break;
                }
                overlapBuilder.prepend(sentence);
            }
            return overlapBuilder.toString();
        }
    }

    int estimateSize(String text) {
        return this.tokenizer != null ? this.tokenizer.estimateTokenCountInText(text) : text.length();
    }

    static TextSegment createSegment(String text, Document document, int index) {
        Metadata metadata = document.getMetadata().copy().put("index", String.valueOf(index));
        return TextSegment.from(text, metadata);
    }
}
